import os
from sales_gpt import SalesGPT
from langchain.chat_models import ChatOpenAI
from langchain.llms.openai import OpenAI, OpenAIChat
from langchain.llms.base import LLM
from typing import Optional, List, Mapping, Any
import requests

os.environ['OPENAI_API_KEY'] = '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' # fill me in

class CustomLLM(LLM):
    model_name = "/mydata2/zengwenjia/project/bloom-1.7b_extract_train"

    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
        prompt_length = len(prompt)
        prompt_dict = {}
        prompt_dict["instruction"] = "你是一个具备广博知识的智能个人助理"
        prompt_dict["input"] = prompt
        # 通过http获取response
        response = requests.get("http://180.184.75.50:7070/get_prompt_response", params=prompt_dict).text
        # inputs = self.pipeline.tokenizer(prompt, return_tensors="pt")
        # input_ids = inputs["input_ids"].to(self.pipeline.device)
        # generation_output = self.pipeline.model.generate(
        #     input_ids=input_ids,
        #     temperature=0.7,
        #     top_p=0.9,
        #     do_sample=True,
        #     num_beams=1,
        #     max_new_tokens=2048,
        #     eos_token_id=self.pipeline.tokenizer.eos_token_id,
        #     pad_token_id=self.pipeline.tokenizer.pad_token_id,
        #     return_dict_in_generate=True,
        #     output_scores=True
        # )
        # s = generation_output.sequences[0]
        # response = self.pipeline.tokenizer.decode(s)
        response = response.replace("</s>", "")
        # response = self.pipeline(prompt, max_new_tokens=num_output)[0]["generated_text"]

        # only return newly generated tokens
        return response

    @property
    def _identifying_params(self) -> Mapping[str, Any]:
        return {"name_of_model": self.model_name}

    @property
    def _llm_type(self) -> str:
        return "custom"

# llm=CustomLLM()
llm = OpenAIChat(temperature=0, model_name="gpt-3.5-turbo")

sales_agent = SalesGPT.from_llm(llm, verbose=True,
                            # salesperson_name="Ted Lasso",
                            # salesperson_role="Sales Representative",
                            # company_name="Sleep Haven",
                            # company_business='''睡眠天堂是一家高级床垫公司，致力于为客户提供尽可能舒适和支撑的睡眠体验。我们提供一系列高品质的床垫、枕头和寝具配件，旨在满足客户的独特需求。'''
                                )

#初始化对话，开场
# sales_agent.seed_agent()
# sales_agent.determine_conversation_stage()

# agent 
# sales_agent.step()
interactive_rounds=1
while True:
    print("interactive_rounds:{}".format(interactive_rounds))
    # user
    if interactive_rounds>1:
        user_input = input('Your response: ')
    else:
        user_input=''
    # sales_agent.human_step(user_input)
    # agent
    # sales_agent.determine_conversation_stage()
    sales_agent.step(session_id='234',query=user_input)
    interactive_rounds+=1